﻿// Copyright (c) 2025 liucl
// Inject is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

#include "process.hpp"

#include <tlhelp32.h>
#include <windows.h>

#include <boost/scope_exit.hpp>

#include <algorithm>
#include <format>
#include <map>
#include <vector>

enum { ID_MENU_INJECT = wxID_HIGHEST + 1, ID_MENU_REFRESH };

struct ProcessTreeItemData : public wxTreeItemData {
  ProcessTreeItemData(DWORD pid, const std::wstring_view &name)
      : pid_(pid), name_(name) {}
  DWORD pid_;
  std::wstring name_;
};

ProcessControl::ProcessControl(wxWindow *parent, wxWindowID id)
    : wxControl(parent, wxID_ANY) {
  auto sizer = new wxBoxSizer(wxVERTICAL);
  processTree_ = new wxTreeCtrl(this, id);
  sizer->Add(processTree_, 1, wxEXPAND);
  SetSizer(sizer);
  processTree_->Bind(wxEVT_TREE_ITEM_RIGHT_CLICK, &ProcessControl::OnRightClick,
                     this);

  EnumProcess();
}

void ProcessControl::OnRightClick(wxTreeEvent &event) {
  auto item = event.GetItem();
  auto data =
      dynamic_cast<ProcessTreeItemData *>(processTree_->GetItemData(item));
  if (data == nullptr) {
    return;
  }
  itemData_ = data;
  wxMenu menu;
  menu.Append(ID_MENU_INJECT, _T("Inject Module"));
  menu.AppendSeparator();
  menu.Append(ID_MENU_REFRESH, _T("Refresh"));
  processTree_->PopupMenu(&menu, event.GetPoint());
  return;
}

void ProcessControl::OnInject(wxCommandEvent &event) {
  assert(itemData_ != nullptr);
  wxCommandEvent cmdEvent(EVT_EVENT_INJECT, processTree_->GetId());
  cmdEvent.SetString(itemData_->name_);
  cmdEvent.SetExtraLong(itemData_->pid_);
  AddPendingEvent(cmdEvent);
}

void ProcessControl::OnRefresh(wxCommandEvent &event) { EnumProcess(); }

void ProcessControl::EnumProcess() {
  processTree_->DeleteAllItems();
  // 枚举进程
  HANDLE hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
  if (hProcessSnap == INVALID_HANDLE_VALUE) {
    return;
  }
  BOOST_SCOPE_EXIT(&hProcessSnap) { CloseHandle(hProcessSnap); }
  BOOST_SCOPE_EXIT_END;

  PROCESSENTRY32 pe32;
  pe32.dwSize = sizeof(PROCESSENTRY32);

  if (!Process32First(hProcessSnap, &pe32)) {
    return;
  }
  std::vector<PROCESSENTRY32> processList;
  do {
    processList.push_back(pe32);
  } while (Process32Next(hProcessSnap, &pe32));

  // 根据父进程ID排序--父进程ID会小于子进程ID
  std::ranges::sort(processList, [](const auto &a, const auto &b) {
    return a.th32ParentProcessID < b.th32ParentProcessID;
  });

  std::map<DWORD, wxTreeItemId> parentMap;
  auto root = processTree_->AddRoot(_T("Process"));
  for (auto pe : processList) {
    auto text =
        std::format(_T("{1} (pid:{0})"), pe.th32ProcessID, pe.szExeFile);
    auto parentId = root;
    if (auto itr = parentMap.find(pe.th32ParentProcessID);
        itr != parentMap.end()) {
      parentId = itr->second;
    }
    auto item = processTree_->AppendItem(
        parentId, text, -1, -1,
        new ProcessTreeItemData(pe.th32ProcessID, pe.szExeFile));
    parentMap[pe.th32ProcessID] = item;
  }
  processTree_->Expand(root);
}
wxDEFINE_EVENT(EVT_EVENT_INJECT, wxCommandEvent);
// clang-format off
wxBEGIN_EVENT_TABLE(ProcessControl, wxControl) 
    EVT_MENU(ID_MENU_INJECT, ProcessControl::OnInject)
    EVT_MENU(ID_MENU_REFRESH, ProcessControl::OnRefresh)
wxEND_EVENT_TABLE()

